import json
import math
from shutil import copyfile

from elasticsearch7 import Elasticsearch

from mlwrite.config import BASE_FOLDER

client = Elasticsearch(hosts=['elastic:1qazxsw2@39.108.14.13:9200', 'elastic:1qazxsw2@120.77.151.145:9200',
                              'elastic:1qazxsw2@39.108.13.45:9200'])


def search_library_text_by_label_id(label_id, base_save_path, contextual=False):
    index = 'library_text'
    count_query = {"query": {"bool": {"should": [{"match_phrase": {"label_id": label_id}}]}}}
    count = client.count(index=index, body=count_query)['count']
    size = 1000
    eponch = math.ceil(count / size)
    query = {
        "from": 0,
        "size": size,
        "track_total_hits": True,
        "query": {
            "bool": {
                "should": [
                    {
                        "match_phrase": {
                            "label_id": label_id
                        }
                    }
                ]
            }
        },
        "sort": [
            {
                "word_num": {
                    "order": "desc"
                }
            }
        ]
    }
    save_path = base_save_path + label_id + '.txt'
    with open(save_path, "w", encoding="utf-8") as corpus_file:
        while eponch != 0:
            print(json.dumps(query))
            res = client.search(index=index, body=json.dumps(query), request_timeout=60)
            hits = res['hits']['hits']
            for text in [hit['_source']['node_text'] for hit in hits]:
                corpus_file.writelines(text.encode('utf-8', 'replace').decode('utf-8'))
                corpus_file.write('\n')
            if len(hits) > 0:
                query['search_after'] = hits[-1]['sort']
            eponch -= 1
    if contextual:
        ty_save_path = BASE_FOLDER + 'contextual/YJtongyzwyj00001/corpus/' + label_id + '.txt'
        copyfile(save_path, ty_save_path)
